#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import os
import shutil
import sys
import threading
from pathlib import Path
from typing import Dict, List, Optional, Sequence

import mimetypes
from PIL import Image, ExifTags, ImageOps
import numpy as np

try:
    from PyQt5 import QtCore, QtGui, QtWidgets
except ImportError as exc:  # pragma: no cover
    raise SystemExit("无法导入 PyQt5，请先执行 pip install PyQt5") from exc

try:
    import face_recognition

    FACE_RECOGNITION_AVAILABLE = True
except ImportError:  # pragma: no cover
    face_recognition = None  # type: ignore
    FACE_RECOGNITION_AVAILABLE = False


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)

try:
    RESAMPLE = Image.Resampling.LANCZOS  # type: ignore[attr-defined]
except AttributeError:  # pragma: no cover
    RESAMPLE = Image.LANCZOS  # type: ignore[attr-defined]

MAX_DETECT_DIMENSION = 1280
FACE_DETECT_UPSAMPLE = 0
HOG_SECOND_PASS_UPSAMPLE = 1
ENABLE_CNN_FALLBACK = True
CNN_FALLBACK_MAX_DIMENSION = 1400
CNN_FALLBACK_UPSAMPLE = 0

class ImageInspector:
    """图片属性检测工具"""
    COMMON_IMAGE_EXTENSIONS = {
        ".jpg",
        ".jpeg",
        ".png",
        ".gif",
        ".bmp",
        ".tiff",
        ".webp",
        ".ico",
        ".jfif",
        ".jpe",
        ".pbm",
        ".pgm",
        ".ppm",
        ".tga",
        ".xbm",
        ".xpm",
        ".psd",
        ".xcf",
        ".ai",
        ".eps",
        ".svg",
        ".heic",
        ".heif",
        ".raw",
        ".arw",
        ".cr2",
        ".cr3",
        ".dcr",
        ".dng",
        ".nef",
        ".orf",
        ".pef",
        ".raf",
        ".rw2",
        ".sr2",
        ".x3f",
        ".3fr",
        ".mef",
        ".mrw",
        ".erf",
    }

    @staticmethod
    def is_image_file(file_path: Path) -> bool:
        """判定文件是否为图片"""
        try:
            mime_type, _ = mimetypes.guess_type(str(file_path))
            if mime_type and mime_type.startswith("image/"):
                return True

            if file_path.suffix.lower() in ImageInspector.COMMON_IMAGE_EXTENSIONS:
                return True

            with Image.open(file_path):
                return True
        except Exception:
            return False

    @staticmethod
    def get_orientation(image_path: Path) -> str:
        """读取图片方向信息"""
        try:
            with Image.open(image_path) as img:
                width, height = img.size
                orientation = 1
                try:
                    exif_dict = img.getexif()
                    if exif_dict:
                        for tag, value in exif_dict.items():
                            if tag in ExifTags.TAGS and ExifTags.TAGS[tag] == "Orientation":
                                orientation = value
                                break
                except Exception:
                    pass

                if orientation in [6, 8]:
                    if height >= width:
                        return "landscape"
                    return "portrait"

                if width >= height:
                    return "landscape"
                return "portrait"
        except Exception:
            return "unknown"


def ensure_directory(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path


def unique_target_path(directory: Path, filename: str) -> Path:
    """生成不重复的目标文件路径"""
    candidate = directory / filename
    if not candidate.exists():
        return candidate

    stem, suffix = os.path.splitext(filename)
    index = 1
    while True:
        candidate = directory / f"{stem}_{index}{suffix}"
        if not candidate.exists():
            return candidate
        index += 1


class ImageSorterWorker(QtCore.QObject):
    """后台线程执行分类任务"""

    log_message = QtCore.pyqtSignal(str)
    progress_changed = QtCore.pyqtSignal(int)
    status_changed = QtCore.pyqtSignal(str)
    finished = QtCore.pyqtSignal(bool, str)

    def __init__(self, folder: Path, mode: str, parent: Optional[QtCore.QObject] = None) -> None:
        super().__init__(parent)
        self.folder = folder
        self.mode = mode
        self._should_stop = threading.Event()

    def stop(self) -> None:
        self._should_stop.set()

    def run(self) -> None:
        try:
            if not self.folder.exists():
                raise FileNotFoundError(f"找不到文件夹: {self.folder}")

            if not os.access(self.folder, os.R_OK | os.W_OK):
                raise PermissionError(f"没有足够的权限访问文件夹: {self.folder}")

            image_files = [
                p for p in self.folder.iterdir() if p.is_file() and ImageInspector.is_image_file(p)
            ]

            if not image_files:
                self.log_message.emit("未找到可处理的图片文件")
                self.finished.emit(False, "未找到图片文件")
                return

            self.status_changed.emit("正在处理图片...")
            self.progress_changed.emit(0)

            if self.mode == "format":
                self._sort_by_format(image_files)
            elif self.mode == "orientation":
                self._sort_by_orientation(image_files)
            elif self.mode == "person":
                self._sort_by_person(image_files)
            else:
                raise ValueError(f"未知的分类模式: {self.mode}")
        except Exception as exc:  # pragma: no cover
            logging.exception("处理过程中发生异常")
            self.log_message.emit(f"发生错误: {exc}")
            self.finished.emit(False, str(exc))
            return
        self.progress_changed.emit(100)
        self.status_changed.emit("分类完成")
        self.log_message.emit("所有任务已完成")
        self.finished.emit(True, "分类完成")

    def _ensure_continue(self) -> bool:
        if self._should_stop.is_set():
            self.log_message.emit("检测到取消信号，提前结束")
            self.finished.emit(False, "用户取消")
            return False
        return True

    def _sort_by_format(self, image_files: Sequence[Path]) -> None:
        total = len(image_files)
        stats: Dict[str, int] = {}

        for index, image_path in enumerate(image_files, start=1):
            if not self._ensure_continue():
                return

            ext = image_path.suffix.lower() or ".unknown"
            stats[ext] = stats.get(ext, 0) + 1
            target = ensure_directory(self.folder / ext.lstrip("."))
            new_path = unique_target_path(target, image_path.name)
            shutil.move(str(image_path), str(new_path))
            self.log_message.emit(f"按格式分类 {image_path.name} -> {new_path.parent.name}")
            self._update_progress(index, total)

        self.log_message.emit("格式分类统计")
        for ext, count in stats.items():
            self.log_message.emit(f"  {ext}: {count} 张")

    def _sort_by_orientation(self, image_files: Sequence[Path]) -> None:
        total = len(image_files)
        stats: Dict[str, int] = {"landscape": 0, "portrait": 0, "unknown": 0}
        labels = {"landscape": "横图", "portrait": "竖图", "unknown": "未知方向"}

        for index, image_path in enumerate(image_files, start=1):
            if not self._ensure_continue():
                return

            orientation = ImageInspector.get_orientation(image_path)
            stats[orientation] = stats.get(orientation, 0) + 1
            target = ensure_directory(self.folder / labels.get(orientation, "未知方向"))
            new_path = unique_target_path(target, image_path.name)
            shutil.move(str(image_path), str(new_path))
            self.log_message.emit(f"按方向分类 {image_path.name} -> {new_path.parent.name}")
            self._update_progress(index, total)

        self.log_message.emit("方向分类统计")
        self.log_message.emit(f"  横图: {stats['landscape']} 张")
        self.log_message.emit(f"  竖图: {stats['portrait']} 张")
        if stats.get("unknown", 0):
            self.log_message.emit(f"  未知方向: {stats['unknown']} 张")

    def _sort_by_person(self, image_files: Sequence[Path]) -> None:
        if not FACE_RECOGNITION_AVAILABLE:
            raise RuntimeError(
                "未安装 face_recognition 库，无法进行人物识别。请执行 pip install face_recognition"
            )

        total = len(image_files)
        recognized_images: Dict[Path, int] = {}
        no_face_images: List[Path] = []

        for index, image_path in enumerate(image_files, start=1):
            if not self._ensure_continue():
                return

            face_locations: List[tuple[int, int, int, int]] = []
            encodings: List[np.ndarray] = []
            try:
                with Image.open(image_path) as pil_image:
                    corrected = ImageOps.exif_transpose(pil_image).convert("RGB")

                working_image = corrected
                max_side = max(working_image.size)
                if max_side > MAX_DETECT_DIMENSION:
                    scale = MAX_DETECT_DIMENSION / float(max_side)
                    new_size = (
                        max(1, int(working_image.width * scale)),
                        max(1, int(working_image.height * scale)),
                    )
                    working_image = working_image.resize(new_size, RESAMPLE)

                image = np.array(working_image)

                face_locations = face_recognition.face_locations(
                    image,
                    number_of_times_to_upsample=FACE_DETECT_UPSAMPLE,
                    model="hog",
                )

                if (
                    not face_locations
                    and HOG_SECOND_PASS_UPSAMPLE is not None
                    and HOG_SECOND_PASS_UPSAMPLE >= 0
                    and HOG_SECOND_PASS_UPSAMPLE != FACE_DETECT_UPSAMPLE
                ):
                    face_locations = face_recognition.face_locations(
                        image,
                        number_of_times_to_upsample=HOG_SECOND_PASS_UPSAMPLE,
                        model="hog",
                    )

                if ENABLE_CNN_FALLBACK and not face_locations:
                    cnn_image = corrected
                    cnn_max_side = max(cnn_image.size)
                    if cnn_max_side > CNN_FALLBACK_MAX_DIMENSION:
                        scale = CNN_FALLBACK_MAX_DIMENSION / float(cnn_max_side)
                        cnn_image = cnn_image.resize((
                            max(1, int(cnn_image.width * scale)),
                            max(1, int(cnn_image.height * scale)),
                        ), RESAMPLE)

                    cnn_array = np.array(cnn_image)
                    cnn_locations = face_recognition.face_locations(
                        cnn_array,
                        number_of_times_to_upsample=CNN_FALLBACK_UPSAMPLE,
                        model="cnn",
                    )

                    if cnn_locations:
                        scale_y = image.shape[0] / float(cnn_array.shape[0])
                        scale_x = image.shape[1] / float(cnn_array.shape[1])

                        def _clamp(value: int, upper: int) -> int:
                            return max(0, min(value, upper))

                        max_y = image.shape[0] - 1
                        max_x = image.shape[1] - 1
                        face_locations = [
                            (
                                _clamp(int(top * scale_y), max_y),
                                _clamp(int(right * scale_x), max_x),
                                _clamp(int(bottom * scale_y), max_y),
                                _clamp(int(left * scale_x), max_x),
                            )
                            for top, right, bottom, left in cnn_locations
                        ]

                if face_locations:
                    encodings = face_recognition.face_encodings(
                        image, known_face_locations=face_locations, num_jitters=1
                    )

            except MemoryError as exc:
                self.log_message.emit(
                    f"读取 {image_path.name} 时出错: 内存不足，已跳过。（{exc}）"
                )
                no_face_images.append(image_path)
                self._update_progress(index, total)
                continue
            except RuntimeError as exc:
                message = str(exc)
                if "bad alloc" in message.lower():
                    self.log_message.emit(
                        f"读取 {image_path.name} 时出错: 内存不足，已跳过。（{exc}）"
                    )
                else:
                    self.log_message.emit(f"读取 {image_path.name} 时出错: {exc}")
                no_face_images.append(image_path)
                self._update_progress(index, total)
                continue
            except Exception as exc:
                self.log_message.emit(f"读取 {image_path.name} 时出错: {exc}")
                no_face_images.append(image_path)
                self._update_progress(index, total)
                continue

            if not encodings:
                no_face_images.append(image_path)
                self.log_message.emit(f"未在 {image_path.name} 中检测到人脸。")
            else:
                recognized_images[image_path] = len(encodings)
            self._update_progress(index, total)

        if not recognized_images and not no_face_images:
            self.log_message.emit("未检测到任何人脸，人物分类终止。")
            return

        people_root = ensure_directory(self.folder / "按人物分类")
        recognized_dir = ensure_directory(people_root / "识别成功")
        for image_path, count in recognized_images.items():
            if not image_path.exists():
                continue
            target_path = unique_target_path(recognized_dir, image_path.name)
            shutil.move(str(image_path), str(target_path))
            self.log_message.emit(
                f"{image_path.name} 检测到 {count} 张人脸，已移动到 识别成功"
            )

        if no_face_images:
            unknown_dir = ensure_directory(people_root / "未识别")
            for image_path in no_face_images:
                if not image_path.exists():
                    continue
                target_path = unique_target_path(unknown_dir, image_path.name)
                shutil.move(str(image_path), str(target_path))
                self.log_message.emit(f"{image_path.name} 移动到 未识别")

        self.log_message.emit(
            f"人物识别统计：成功 {len(recognized_images)} 张，未识别 {len(no_face_images)} 张"
        )


    def _update_progress(self, current: int, total: int) -> None:
        progress = int((current / total) * 100)
        self.progress_changed.emit(progress)
        self.status_changed.emit(f"正在处理 {current}/{total}")


class ImageSorterWindow(QtWidgets.QMainWindow):
    """PyQt5 主界面"""

    def __init__(self) -> None:
        super().__init__()
        self.setWindowTitle("图片分类助手")
        self.resize(860, 640)

        self._thread: Optional[QtCore.QThread] = None
        self._worker: Optional[ImageSorterWorker] = None

        self._build_ui()
        self._connect_signals()

    def _build_ui(self) -> None:
        central = QtWidgets.QWidget(self)
        self.setCentralWidget(central)

        layout = QtWidgets.QVBoxLayout(central)
        layout.setContentsMargins(16, 16, 16, 16)
        layout.setSpacing(12)

        title = QtWidgets.QLabel("图片智能分类")
        title_font = QtGui.QFont()
        title_font.setPointSize(18)
        title_font.setBold(True)
        title.setFont(title_font)
        layout.addWidget(title)

        folder_group = QtWidgets.QGroupBox("选择文件夹")
        folder_layout = QtWidgets.QHBoxLayout(folder_group)
        self.folder_edit = QtWidgets.QLineEdit()
        self.folder_edit.setPlaceholderText("请选择需要分类的图片所在文件夹")
        browse_button = QtWidgets.QPushButton("浏览")
        folder_layout.addWidget(self.folder_edit)
        folder_layout.addWidget(browse_button)
        layout.addWidget(folder_group)

        mode_group = QtWidgets.QGroupBox("分类模式")
        mode_layout = QtWidgets.QHBoxLayout(mode_group)
        self.mode_combo = QtWidgets.QComboBox()
        self.mode_combo.addItem("按文件格式", "format")
        self.mode_combo.addItem("按图片方向", "orientation")
        self.mode_combo.addItem("按人物分类", "person")
        mode_layout.addWidget(self.mode_combo)
        layout.addWidget(mode_group)

        log_group = QtWidgets.QGroupBox("运行日志")
        log_layout = QtWidgets.QVBoxLayout(log_group)
        self.log_view = QtWidgets.QTextEdit()
        self.log_view.setReadOnly(True)
        log_layout.addWidget(self.log_view)
        layout.addWidget(log_group, 1)

        progress_group = QtWidgets.QGroupBox("处理进度")
        progress_layout = QtWidgets.QVBoxLayout(progress_group)
        self.progress_bar = QtWidgets.QProgressBar()
        self.status_label = QtWidgets.QLabel("等待开始")
        progress_layout.addWidget(self.progress_bar)
        progress_layout.addWidget(self.status_label)
        layout.addWidget(progress_group)

        button_layout = QtWidgets.QHBoxLayout()
        button_layout.addStretch(1)
        self.start_button = QtWidgets.QPushButton("开始分类")
        self.clear_button = QtWidgets.QPushButton("清除日志")
        button_layout.addWidget(self.start_button)
        button_layout.addWidget(self.clear_button)
        layout.addLayout(button_layout)

        self._browse_button = browse_button

    def _connect_signals(self) -> None:
        self._browse_button.clicked.connect(self._choose_folder)
        self.start_button.clicked.connect(self._start_sorting)
        self.clear_button.clicked.connect(self.log_view.clear)

    def _choose_folder(self) -> None:
        folder = QtWidgets.QFileDialog.getExistingDirectory(self, "选择图片文件夹", str(Path.home()))
        if folder:
            self.folder_edit.setText(folder)
            self._append_log(f"选定文件夹: {folder}")

    def _start_sorting(self) -> None:
        if self._thread is not None:
            QtWidgets.QMessageBox.warning(self, "正在运行", "任务执行中，请稍候")
            return

        folder = Path(self.folder_edit.text().strip())
        if not folder:
            QtWidgets.QMessageBox.warning(self, "提示", "请先选择需要分类的文件夹")
            return

        mode = self.mode_combo.currentData()
        self.progress_bar.setValue(0)
        self.status_label.setText("准备开始")
        self._append_log(f"开始分类：{folder}（模式：{self.mode_combo.currentText()}）")

        self._thread = QtCore.QThread(self)
        self._worker = ImageSorterWorker(folder, mode)
        self._worker.moveToThread(self._thread)

        self._thread.started.connect(self._worker.run)
        self._worker.log_message.connect(self._append_log)
        self._worker.progress_changed.connect(self.progress_bar.setValue)
        self._worker.status_changed.connect(self.status_label.setText)
        self._worker.finished.connect(self._on_worker_finished)

        self._thread.start()
        self.start_button.setEnabled(False)

    def _append_log(self, message: str) -> None:
        self.log_view.append(message)
        self.log_view.verticalScrollBar().setValue(self.log_view.verticalScrollBar().maximum())
        logging.info(message)

    def closeEvent(self, event: QtGui.QCloseEvent) -> None:  # noqa: N802
        if self._worker is not None:
            self._worker.stop()
        if self._thread is not None:
            self._thread.quit()
            self._thread.wait(2000)
        super().closeEvent(event)

    def _on_worker_finished(self, success: bool, message: str) -> None:
        if self._thread is None or self._worker is None:
            return

        self._thread.quit()
        self._thread.wait()
        self._thread = None
        self._worker = None
        self.start_button.setEnabled(True)

        if success:
            QtWidgets.QMessageBox.information(self, "完成", message)
        else:
            QtWidgets.QMessageBox.warning(self, "提示", message)


def main() -> None:
    app = QtWidgets.QApplication(sys.argv)
    window = ImageSorterWindow()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()


